library(PoissonBinomial)
library(FNN)
library(GoFKernel)
library(LaplacesDemon)

# preliminaries
set.seed(2022)
n <- 1000

# statistic to report (options are 'correlation', 'rSquared', and 'rmse'. 
# 'correlation' is used in the paper) 
mode <- 'correlation' 

# list of the six data-generating processes 
dgps <- c('uniform', 'closeToZero', 'closeToOne', 'extremal', 'central', 'bimodal')

############################################
##            helper functions            ##
############################################

# logit swing function
logitSwing <- function(probs, D, lwr = 0, upr = 5e4, eps = 1e-8) {
  
  # compute logit swing
  while(upr - lwr > eps) {
    mid <- (lwr + upr)/2
    
    if(sum(1/(1 + (1-probs)/probs*mid)) < D)
      upr <- mid
    else
      lwr <- mid
  }

  return(1/(1 + (1-probs)/probs*mid))
}

# close to zero quantile function
extremalQuantile <- inverse(function(x) {
    1/2*(pbeta(x, 0.1, 3) + pbeta(x, 3, 0.1))
  }, lower = 0, upper = 1)
extrema <- Vectorize(extremalQuantile)

# bimodal quantile function
bimodalQuantile <- inverse(function(x) {
  1/2*(pbeta(x, 3, 10) + pbeta(x, 10, 3))
}, lower = 0, upper = 1)
bimodal <- Vectorize(bimodalQuantile)

# data generation function
dataGenerate <- function(x, dist) {
  
  # generate the data
  if(dist == 'uniform')
    p <- qunif(x)
  else if(dist == 'closeToZero')
    p <- qbeta(x, 0.1, 3)
  else if(dist == 'closeToOne')
    p <- qbeta(x, 3, 0.1)
  else if(dist == 'extremal')
    p <- extrema(x)
  else if(dist == 'central')
    p <- qbeta(x, 3, 3)
  else if(dist == 'bimodal')
    p <- bimodal(x)

  return(p)
}

############################################
##             calling script             ##
############################################

# iterate through every combination of true and initial prediction distributions
results <- sapply(dgps, FUN = function(dgpTrue) {
  sapply(dgps, FUN = function(dgpEst) {
    
    # source the data such that the ordering is unchanged 
    p <- runif(n)
    trueScores <- dataGenerate(p, dgpTrue)
    estScores <- dataGenerate(p, dgpEst)
    
    # compute the logit swing
    D <- sum(outcomes <- as.numeric(runif(n) < trueScores))
    updatedScores <- logitSwing(estScores, D)
    
    # report the result statistics
    if(mode == 'correlation') { # these are the results reported in the paper, Table 3
      cor.prior <- cor(trueScores, estScores)
      cor.posterior <-  cor(trueScores, updatedScores)
      (cor.posterior - cor.prior)/cor.prior
      
    } else if(mode == 'rmse') {
      rmse.prior <- sqrt(mean((trueScores - estScores)^2))
      rmse.posterior <- sqrt(mean((trueScores - updatedScores)^2))
      rmse.posterior - rmse.prior
      
    } else if(mode == 'rSquared') {
      rSq.prior <- 1 - sum((trueScores - estScores)^2)/sum((trueScores - mean(trueScores))^2)
      rSq.posterior <- 1 - sum((trueScores - updatedScores)^2)/sum((trueScores - mean(trueScores))^2)
      rSq.posterior - rSq.prior
    }
    
  })
})

# print the results
Table3 <- round(results, 3)
print(Table3)
